import numpy as np
import random

import torch
from torch import nn
import torch.optim as optim

from PPINN.trainer import PINN, PINN2D, PINN3D
from Plot.PlotSolution import plot_solution3D, plot_solution2D

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

N_u=500
N_f=5000
mesh_path = 'data/PINN2D/mesh.npz'

boundary_y = np.array([[0.25, 0.04, 0.04, -0.04, -0.04, -0.45, -0.45, 0.25]])
boundary_z = np.array([[0.05, 0.05, 0.25, 0.25, 0.05, 0.05, -0.05, -0.05]])
mask =  np.vstack((boundary_y, boundary_z))
def generate_data():
    x_upper = np.ones((N_u // 4, 1), dtype=float)
    x_lower = np.ones((N_u // 4, 1), dtype=float) * (-1)
    t_zero = np.zeros((N_u // 2, 1), dtype=float)

    t_upper = np.random.rand(N_u // 4, 1)
    t_lower = np.random.rand(N_u // 4, 1)
    x_zero = (-1) + np.random.rand(N_u // 2, 1) * (1 - (-1))

    X_upper = np.hstack((x_upper, t_upper))
    X_lower = np.hstack((x_lower, t_lower)) 
    X_zero = np.hstack((x_zero, t_zero))   

    X_u_train = np.vstack((X_upper, X_lower, X_zero))

    index = np.arange(0, N_u)
    np.random.shuffle(index)
    X_u_train = X_u_train[index, :]

    X_f_train = np.zeros((N_f, 2), dtype=float)
    for row in range(N_f):
        x = random.uniform(-1, 1)  
        t = random.uniform(0, 1)  
        X_f_train[row, 0] = x
        X_f_train[row, 1] = t

    X_f_train = np.vstack((X_f_train, X_u_train))

    return(X_f_train)

def sampling_from_msh(mesh_path, N_u=None, N_f=None):
    loaded_data = np.load(mesh_path)
    boundary_number = 0
    interior_number = 0
    index = {}
    coord = {}
    for key, value in loaded_data.items():
        coord[key] = []
        if key == 'interior':
            interior_number = len(value)
        else:
            index[key] = np.zeros((1, 3), dtype=int) # 起始、终止、数量
            index[key][0, 0] = boundary_number
            index[key][0, 1] = boundary_number + len(value)  - 1          
            boundary_number = boundary_number + len(value)
    if N_u is None or N_f is None:
        boundary_indexs = random.sample(range(boundary_number), boundary_number)
        interior_indexs = random.sample(range(interior_number), interior_number)
    else:
        boundary_indexs = random.sample(range(boundary_number), N_u)
        interior_indexs = random.sample(range(interior_number), N_f)

    for i in range(len(boundary_indexs)):
        for key, value in index.items():
            if boundary_indexs[i] <= value[0, 1] and boundary_indexs[i] >= value[0, 0]:
                value[0, 2] = value[0, 2] + 1
                location = int(boundary_indexs[i] - value[0, 0])
                coord[key].append(loaded_data[key][location])
                break

    coord['interior'] = loaded_data['interior'][interior_indexs]

    return coord
if __name__ == "__main__":
    X_f_train = sampling_from_msh(mesh_path, N_u, N_f)
    for key, value in X_f_train.items():
        X_f_train[key] = torch.tensor(value, dtype=torch.float32, device=device)
    # X_f_train = generate_data()
    # X_f_train = torch.tensor(X_f_train, dtype=torch.float32, device=device)

    pin = PINN2D(X_f_train, N_u, N_f, device)
    pin.train()

    X_f_train = sampling_from_msh(mesh_path)
    plot_solution2D(pin, X_f_train)